/*
 * Copyright 2016 The Cartographer Authors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "slam/common/time.h"
#include "slam_ros/ros_log_sink.h"
#include "slam_ros/time_conversion.h"
#include "gflags/gflags.h"
#include "agv_msgs/msg/odometry_lite.hpp"
#include "rclcpp/rclcpp.hpp"
#include "rosbag2_cpp/reader.hpp"
#include "sensor_msgs/msg/imu.hpp"
#include "sensor_msgs/msg/laser_scan.hpp"
#include "sensor_msgs/msg/multi_echo_laser_scan.hpp"
#include "sensor_msgs/msg/point_cloud2.hpp"
#include "tf2_msgs/msg/tf_message.hpp"
#include "rclcpp/serialization.hpp"

DEFINE_string(bag_filename, "", "Bag to publish.");

const int kQueueSize = 1;

// 通用消息处理模板
template <typename T>
void PublishWithModifiedTimestamp(
    const std::shared_ptr<T>& message,
    const typename rclcpp::Publisher<T>::SharedPtr& publisher,
    const rclcpp::Duration& bag_to_current) {
  // Convert duration to seconds and nanoseconds
  auto duration_ns = bag_to_current.nanoseconds();
  rclcpp::Time new_stamp(message->header.stamp.sec, message->header.stamp.nanosec);
  new_stamp = new_stamp + rclcpp::Duration::from_nanoseconds(duration_ns);
  message->header.stamp.sec = new_stamp.seconds();
  message->header.stamp.nanosec = new_stamp.nanoseconds() % 1000000000;
  publisher->publish(*message);
}

// TFMessage特化版本
template <>
void PublishWithModifiedTimestamp<tf2_msgs::msg::TFMessage>(
    const std::shared_ptr<tf2_msgs::msg::TFMessage>& message,
    const typename rclcpp::Publisher<tf2_msgs::msg::TFMessage>::SharedPtr& publisher,
    const rclcpp::Duration& bag_to_current) {
  auto duration_ns = bag_to_current.nanoseconds();
  for (auto& transform : message->transforms) {
    rclcpp::Time new_stamp(transform.header.stamp.sec, transform.header.stamp.nanosec);
    new_stamp = new_stamp + rclcpp::Duration::from_nanoseconds(duration_ns);
    transform.header.stamp.sec = new_stamp.seconds();
    transform.header.stamp.nanosec = new_stamp.nanoseconds() % 1000000000;
  }
  publisher->publish(*message);
}

int main(int argc, char** argv) {
  google::InitGoogleLogging(argv[0]);
  google::SetUsageMessage(
      "\n\n"
      "This replays and publishes messages from a given bag file, modifying "
      "their header timestamps to match current ROS time.\n\n");
  google::ParseCommandLineFlags(&argc, &argv, true);
  CHECK(!FLAGS_bag_filename.empty()) << "-bag_filename is missing.";

  rclcpp::init(argc, argv);
  auto node = std::make_shared<rclcpp::Node>("rosbag_publisher");

  // 初始化序列化对象
  rclcpp::Serialization<sensor_msgs::msg::PointCloud2> pc2_serialization;
  rclcpp::Serialization<sensor_msgs::msg::Imu> imu_serialization;
  rclcpp::Serialization<sensor_msgs::msg::LaserScan> scan_serialization;
  rclcpp::Serialization<sensor_msgs::msg::MultiEchoLaserScan> multi_echo_serialization;
  rclcpp::Serialization<agv_msgs::msg::OdometryLite> odom_serialization;
  rclcpp::Serialization<tf2_msgs::msg::TFMessage> tf_serialization;

  rosbag2_cpp::Reader reader;
  reader.open(FLAGS_bag_filename);

  std::map<std::string, rclcpp::PublisherBase::SharedPtr> topic_to_publisher;
  auto metadata = reader.get_metadata();

  // 初始化发布器
  for (const auto& topic_info : metadata.topics_with_message_count) {
    const std::string& topic = topic_info.topic_metadata.name;
    if (topic_to_publisher.count(topic) == 0) {
      if (topic == "/point_cloud") {
        topic_to_publisher[topic] = node->create_publisher<sensor_msgs::msg::PointCloud2>(
            topic, kQueueSize);
      } else if (topic == "/imu") {
        topic_to_publisher[topic] = node->create_publisher<sensor_msgs::msg::Imu>(
            topic, kQueueSize);
      } else if (topic == "/scan") {
        topic_to_publisher[topic] = node->create_publisher<sensor_msgs::msg::LaserScan>(
            topic, kQueueSize);
      } else if (topic == "/multi_echo_scan") {
        topic_to_publisher[topic] = node->create_publisher<sensor_msgs::msg::MultiEchoLaserScan>(
            topic, kQueueSize);
      } else if (topic == "/odometry_lite") {
        topic_to_publisher[topic] = node->create_publisher<agv_msgs::msg::OdometryLite>(
            topic, kQueueSize);
      } else if (topic == "/tf") {
        topic_to_publisher[topic] = node->create_publisher<tf2_msgs::msg::TFMessage>(
            topic, kQueueSize);
      }
    }
  }

  const auto current_start = node->now();
  const auto bag_start = rclcpp::Time(metadata.starting_time.time_since_epoch().count());
  const auto bag_to_current = current_start - bag_start;

  while (reader.has_next() && rclcpp::ok()) {
    auto serialized_msg = reader.read_next();
    const auto message_time = rclcpp::Time(serialized_msg->time_stamp);
    const auto after_bag_start = message_time - bag_start;
    const auto planned_publish_time = current_start + after_bag_start;

    const std::string& topic = serialized_msg->topic_name;

    try {
      if (topic == "/point_cloud") {
        auto msg = std::make_shared<sensor_msgs::msg::PointCloud2>();
        rclcpp::SerializedMessage extracted_msg(*serialized_msg->serialized_data);
        pc2_serialization.deserialize_message(&extracted_msg, msg.get());
        auto publisher = std::dynamic_pointer_cast<rclcpp::Publisher<sensor_msgs::msg::PointCloud2>>(
            topic_to_publisher[topic]);
        PublishWithModifiedTimestamp(msg, publisher, bag_to_current);
      } 
      else if (topic == "/imu") {
        auto msg = std::make_shared<sensor_msgs::msg::Imu>();
        rclcpp::SerializedMessage extracted_msg(*serialized_msg->serialized_data);
        imu_serialization.deserialize_message(&extracted_msg, msg.get());
        auto publisher = std::dynamic_pointer_cast<rclcpp::Publisher<sensor_msgs::msg::Imu>>(
            topic_to_publisher[topic]);
        PublishWithModifiedTimestamp(msg, publisher, bag_to_current);
      }
      else if (topic == "/scan") {
        auto msg = std::make_shared<sensor_msgs::msg::LaserScan>();
        rclcpp::SerializedMessage extracted_msg(*serialized_msg->serialized_data);
        scan_serialization.deserialize_message(&extracted_msg, msg.get());
        auto publisher = std::dynamic_pointer_cast<rclcpp::Publisher<sensor_msgs::msg::LaserScan>>(
            topic_to_publisher[topic]);
        PublishWithModifiedTimestamp(msg, publisher, bag_to_current);
      }
      else if (topic == "/multi_echo_scan") {
        auto msg = std::make_shared<sensor_msgs::msg::MultiEchoLaserScan>();
        rclcpp::SerializedMessage extracted_msg(*serialized_msg->serialized_data);
        multi_echo_serialization.deserialize_message(&extracted_msg, msg.get());
        auto publisher = std::dynamic_pointer_cast<rclcpp::Publisher<sensor_msgs::msg::MultiEchoLaserScan>>(
            topic_to_publisher[topic]);
        PublishWithModifiedTimestamp(msg, publisher, bag_to_current);
      }
      else if (topic == "/odometry_lite") {
        auto msg = std::make_shared<agv_msgs::msg::OdometryLite>();
        rclcpp::SerializedMessage extracted_msg(*serialized_msg->serialized_data);
        odom_serialization.deserialize_message(&extracted_msg, msg.get());
        auto publisher = std::dynamic_pointer_cast<rclcpp::Publisher<agv_msgs::msg::OdometryLite>>(
            topic_to_publisher[topic]);
        PublishWithModifiedTimestamp(msg, publisher, bag_to_current);
      }
      else if (topic == "/tf") {
        auto msg = std::make_shared<tf2_msgs::msg::TFMessage>();
        rclcpp::SerializedMessage extracted_msg(*serialized_msg->serialized_data);
        tf_serialization.deserialize_message(&extracted_msg, msg.get());
        auto publisher = std::dynamic_pointer_cast<rclcpp::Publisher<tf2_msgs::msg::TFMessage>>(
            topic_to_publisher[topic]);
        PublishWithModifiedTimestamp(msg, publisher, bag_to_current);
      }
    } catch (const std::exception& e) {
      RCLCPP_ERROR(node->get_logger(), "Failed to process %s: %s", 
                 topic.c_str(), e.what());
    }

    const auto current_time = node->now();
    double simulation_delay = slam::common::ToSeconds(
        slam_ros::FromRos(current_time) -
        slam_ros::FromRos(planned_publish_time));
    if (std::abs(simulation_delay) > 0.001) {
      RCLCPP_WARN(node->get_logger(), 
                 "Playback delayed by %f s. planned_publish_time: %f current_time: %f",
                 simulation_delay,
                 planned_publish_time.seconds(),
                 current_time.seconds());
    }
  }

  rclcpp::shutdown();
  return 0;
}